import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt

import input, preprocess

# 计算各国的项目相关性
def correlate(medal_counts):
    columns_to_analyze = ['Aquatics_Ratio', 'Ball_Sports_Ratio', 'Equestrian_Ratio', 
                    'Martial_Arts_Ratio', 'Extreme_Sports_Ratio', 'Other_Sports_Ratio']
    # 初始化一个空的列表来存储每个国家的相关性矩阵
    correlation_results = []
    # 对每个国家分别计算相关性
    for country, group in medal_counts.groupby('Country'):
    # 选择需要分析的列
        group_numeric = group[columns_to_analyze + ['Total']]
        # 填充缺失值，例如用0填充
        group_numeric = group_numeric.fillna(0)
        # 计算相关性矩阵
        correlation_matrix = group_numeric.corr()    
        # 提取与 Total 的相关性
        total_correlations = correlation_matrix['Total'].drop('Total')
        # 保存相关性结果
        correlation_results.append({
            'Country': country,
            **total_correlations.round(2).to_dict()
         })
    correlation_df = pd.DataFrame(correlation_results)
    # 打印每个国家的相关性矩阵
    # print(correlation_df)
    # correlation_df.to_csv('./correlation_results.csv', index=False)

    return correlation_df


# 寻找特定国家并画出热力图
def draw_picture(correlation_df, country):
    country_data = correlation_df[correlation_df['Country'] == country]
    
    # 如果找到该国家的数据，则绘制热力图
    if not country_data.empty:
        # 将相关性结果转换为宽格式
        country_corr = country_data.set_index('Country').T

        # 绘制热力图
        plt.figure(figsize=(10, 6))  # 调整热力图的大小
        sns.heatmap(country_corr, annot=True, cmap='coolwarm', fmt=".2f", linewidths=.5, annot_kws={"size": 12})
        plt.title(f"Correlation of Sports Ratios with Total Medals for {country}", fontsize=14)
        plt.xticks(rotation=0, ha='right', fontsize=9)  # 调整X轴标签的字体大小和旋转角度
        plt.yticks(rotation=-45, fontsize=9)  # 调整Y轴标签的字体大小
        plt.show()
    else:
        print(f"Country {country} not found in the data.")



# 主程序
if __name__ == "__main__":
    file_path = "./data/"  # 替换为你的文件路径
    dataframes = input.load_data(file_path)
    medal_counts = preprocess.preprocess_data(dataframes)

    # 计算相关性
    correlation_df = correlate(medal_counts)

    # 可视化相关性结果, 选择一个国家作为示例
    country = 'United States'
    draw_picture(correlation_df, country)